#!/usr/bin/env python3
"""
IMU → 世界系点云 逐步累积显示 + IMU 轨迹（红点）
python3.12 bag2world_accum.py [--bag xxx.db3] [--bunch 20]
"""
import argparse, numpy as np, matplotlib.pyplot as plt
from functools import partial
from matplotlib import animation
from rosbag2_py import SequentialReader, StorageOptions, ConverterOptions, StorageFilter
from sensor_msgs.msg import PointCloud2, Imu
from rclpy.serialization import deserialize_message

# ---------- 工具 ----------
def quat2R(qw, qx, qy, qz):
    R = np.array([
        [1 - 2 * (qy ** 2 + qz ** 2), 2 * (qx * qy - qw * qz), 2 * (qx * qz + qw * qy)],
        [2 * (qx * qy + qw * qz), 1 - 2 * (qx ** 2 + qz ** 2), 2 * (qy * qz - qw * qx)],
        [2 * (qx * qz - qw * qy), 2 * (qy * qz + qw * qx), 1 - 2 * (qx ** 2 + qy ** 2)]
    ], dtype=np.float32)
    return R

def pc2_to_xyz(pc: PointCloud2):
    arr = np.frombuffer(pc.data, np.uint8).reshape(-1, pc.point_step)
    xyz = arr[:, :12].copy().view(np.float32).reshape(-1, 3)
    return xyz[np.isfinite(xyz).all(axis=1)]

# ---------- 加载 + 累积 ----------
def load_accumulated(bag_path, bunch_size):
    reader = SequentialReader()
    reader.open(StorageOptions(uri=bag_path, storage_id='sqlite3'),
                ConverterOptions('', ''))
    reader.set_filter(StorageFilter(topics=['/unilidar/imu', '/unilidar/cloud']))

    world_points, imu_track = [], []          # 累积点云 & IMU 轨迹
    latest_R = np.eye(3, dtype=np.float32)
    latest_t = np.zeros(3, dtype=np.float32)  # IMU 世界位置（初始原点）
    bunches = []                              # 每步要追加的新点云/轨迹
    cache_pc, cache_imu = [], []              # 帧缓存

    while reader.has_next():
        topic, data, t = reader.read_next()
        if topic == '/unilidar/imu':
            imu = deserialize_message(data, Imu)
            q = imu.orientation
            latest_R = quat2R(q.w, q.x, q.y, q.z)
            # 假设 IMU 在世界系原点移动，可换成积分速度或位移
            latest_t = np.zeros(3, dtype=np.float32)
            cache_imu.append(latest_t.copy())
        elif topic == '/unilidar/cloud':
            pc = deserialize_message(data, PointCloud2)
            pts_imu = pc2_to_xyz(pc)
            pts_world = (latest_R @ pts_imu.T).T + latest_t
            cache_pc.append(pts_world)

            if len(cache_pc) == bunch_size:
                # 新点云 & 新 IMU 轨迹
                new_pcd = np.vstack(cache_pc)
                new_imu = np.vstack(cache_imu) if cache_imu else np.empty((0, 3))
                world_points.append(new_pcd)
                imu_track.append(new_imu)
                bunches.append((new_pcd, new_imu))
                cache_pc, cache_imu = [], []
    # 尾部
    if cache_pc:
        new_pcd = np.vstack(cache_pc)
        new_imu = np.vstack(cache_imu) if cache_imu else np.empty((0, 3))
        world_points.append(new_pcd)
        imu_track.append(new_imu)
        bunches.append((new_pcd, new_imu))

    # 一次性计算全局坐标轴范围
    all_pts = np.vstack(world_points) if world_points else np.empty((0, 3))
    return bunches, all_pts

# ---------- 动画 ----------
def animate(i, bunches, global_min, global_max):
    new_pcd, new_imu = bunches[i]
    print(f'step {i+1}/{len(bunches)}  points={len(new_pcd)} imu_pts={len(new_imu)}')
    # 追加新点云（青色小点）
    # scat_pcd._offsets3d = (np.concatenate([scat_pcd._offsets3d[0], new_pcd[:, 0]]),
    #                        np.concatenate([scat_pcd._offsets3d[1], new_pcd[:, 1]]),
    #                        np.concatenate([scat_pcd._offsets3d[2], new_pcd[:, 2]]))
    # 追加 IMU 轨迹（红色大点）
    if new_imu.shape[0] > 0:
        scat_imu._offsets3d = (np.concatenate([scat_imu._offsets3d[0], new_imu[:, 0]]),
                               np.concatenate([scat_imu._offsets3d[1], new_imu[:, 1]]),
                               np.concatenate([scat_imu._offsets3d[2], new_imu[:, 2]]))
    ax.set_title(f'step {i+1}/{len(bunches)}  total_pts={len(scat_pcd._offsets3d[0])}')
    return scat_pcd, scat_imu

# ---------- main ----------
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--bag', default='/gitlab/pcd_analysis/demo_data/unilidar-2023-09-22-12-42-04.db3',
                        help='ros2 bag db3 path')
    parser.add_argument('--bunch', type=int, default=20, help='frames per step')
    parser.add_argument('--interval', type=int, default=50, help='ms per step')
    args = parser.parse_args()

    print('读取 bag 并累积世界系点云 + IMU 轨迹 …')
    bunches, global_pts = load_accumulated(args.bag, args.bunch)
    print(f'共 {len(bunches)} 步，总点数 {len(global_pts)}')

    fig = plt.figure(figsize=(8, 6))
    global ax, scat_pcd, scat_imu
    ax = fig.add_subplot(111, projection='3d')
    # 初始空散点
    scat_pcd = ax.scatter([], [], [], s=1, c='cyan', label='point cloud')
    scat_imu = ax.scatter([], [], [], s=10, c='red', label='IMU track')
    ax.set_box_aspect([1, 1, 1])
    ax.set_xlabel('X world')
    ax.set_ylabel('Y world')
    ax.set_zlabel('Z world')
    ax.legend()

    # 一次性设定坐标轴范围
    margin = 1.0
    mins, maxs = global_pts.min(axis=0) - margin, global_pts.max(axis=0) + margin
    ax.set_xlim(mins[0], maxs[0])
    ax.set_ylim(mins[1], maxs[1])
    ax.set_zlim(mins[2], maxs[2])

    ani = animation.FuncAnimation(fig, partial(animate, bunches=bunches,
                                               global_min=mins, global_max=maxs),
                                  frames=len(bunches), interval=args.interval,
                                  blit=False, repeat=True)
    plt.show()

if __name__ == '__main__':
    main()